{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.metrics import confusion_matrix\n",
    "import pickle\n",
    "import random\n",
    "from copy import deepcopy\n",
    "import itertools\n",
    "import datetime\n",
    "import time\n",
    "import json\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "from torch.optim.lr_scheduler import LambdaLR\n",
    "from sympy import symbols\n",
    "import pprint as pp\n",
    "\n",
    "import sys, os\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "from AI_physicist.theory_learning.models import Loss_Fun_Cumu\n",
    "from AI_physicist.theory_learning.theory_model import Theory_Training, get_loss, get_best_model_idx, get_preds_valid, get_valid_prob\n",
    "from AI_physicist.theory_learning.theory_hub import Theory_Hub, load_model_dict_at_theory_hub\n",
    "from AI_physicist.theory_learning.util_theory import plot_theories, plot3D, process_loss, plot_loss_record, to_one_hot, plot_indi_domain, get_mystery, to_pixels, standardize_symbolic_expression, check_expression_matching\n",
    "from AI_physicist.theory_learning.util_theory import get_epochs, export_csv_with_domain_prediction, extract_target_expression, get_coeff_learned_list, load_theory_training, get_dataset, get_piecewise_dataset\n",
    "from AI_physicist.settings.filepath import theory_PATH\n",
    "from AI_physicist.settings.global_param import COLOR_LIST\n",
    "from AI_physicist.pytorch_net.util import Loss_Fun, make_dir, Loss_with_uncertainty, Early_Stopping, record_data, plot_matrices, get_args, serialize, to_string, filter_filename, to_np_array, to_Variable\n",
    "from AI_physicist.pytorch_net.net import Conv_Autoencoder, load_model_dict, train_simple\n",
    "\n",
    "try:\n",
    "    %matplotlib inline\n",
    "    isplot = True\n",
    "except:\n",
    "    import matplotlib\n",
    "    matplotlib.use('Agg')\n",
    "    isplot = False\n",
    "is_cuda = torch.cuda.is_available()\n",
    "standardize = standardize_symbolic_expression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset loading and important settings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "########################################\n",
    "# Setting up path to dataset files. The data files are csv_dirname + env_name + \".csv\", \n",
    "# where env_name are elements in the csv_filename_list. For each csv file, the first\n",
    "# (num_output_dims * num_input_steps) columns are the past (E.g., if num_output_dims=2, then\n",
    "# they are arranged as (x_{t-num_input_steps}, y_{t-num_input_steps}, ... x_{t-1}, y_{t-1}) ).\n",
    "# The next num_output_dims columns are the target prediction for the future.\n",
    "# If is_classified = True, the last column in the csv should provide the true_domain id for evaluating \n",
    "# whether the domain prediction is correct (not used for training)\n",
    "########################################\n",
    "csv_dirname = \"../datasets/GROUND_TRUTH/\"  # Path for the dataset file\n",
    "is_classified = True             # If True, the last column in the csv file should provide the true_domain id for evaluation.\n",
    "csv_filename_list = get_mystery(50000, range(4,7), range(1, 6), is_classified) + get_mystery(50000, [20], range(1, 6), is_classified) + get_mystery(50000, range(7, 11), range(1, 6), is_classified)\n",
    "print(\"Environments:\\n\", csv_filename_list)\n",
    "num_output_dims = 2               # It sets the dimension of output\n",
    "num_input_steps = 2               # It sets the number of steps for the input\n",
    "\n",
    "########################################\n",
    "# Other important settings:\n",
    "########################################\n",
    "exp_mode = \"continuous\"           # Choose from \"continuous\" (full AI Physicist), \"newb\" (newborn) and \"base\" (baseline)\n",
    "forward_steps = 1                 # Number of forward steps to predict\n",
    "data_format = \"states\"            # Choose from \"states\" or \"images\"\n",
    "pred_nets_activation = \"linear\"   # Activation for the prediction function f. Choose from \"linear\", \"leakyRelu\"\n",
    "num_layers = 3                    # Number of layers for the prediction function f.\n",
    "\n",
    "num_theories_init = 4             # Number of theories to start with.\n",
    "add_theory_loss_threshold = 2e-6  # MSE threshold for individual data points above which to add a new theory to fit.\n",
    "add_theory_criteria = (\"loss_with_domain\", 0)  # Criteria and threshold of loss increase to determine whether to accept adding the theory.\n",
    "add_theory_quota = 1              # maximum number of theories to add at each phase.\n",
    "add_theory_limit = None           # maximum allowed number of theories. If None, do not set limit.\n",
    "\n",
    "is_Lagrangian = False             # If True, learn the Lagrangian. If False, learn Equation of Motion (EOM).\n",
    "load_previous = True              # Whether to load previously trained instances on"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings up:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "../data/theory/exp1.0_11-3/file_continuous_num_4_pred_8_linear_dom_8_onehot_mse_1e-07_sim_DLs-0-3-relative_optim_adam-0.005_adam-0.001_reg_1e-08_1e-05_batch_10000_core_DLs_order_-1_lossd_None_False_infl_False_#_20000_mul_5000_MDL_both_2D_3L_id_0_0.p\n"
     ]
    }
   ],
   "source": [
    "# Other settings:\n",
    "exp_id = \"exp1.0\"\n",
    "env_source = \"file\"\n",
    "pred_nets_neurons = 8\n",
    "domain_net_neurons = 8\n",
    "domain_pred_mode = \"onehot\"\n",
    "mse_amp = 1e-7\n",
    "scheduler_settings = (\"ReduceLROnPlateau\", 40, 0.1)   # Settings for the learning rate scheduler\n",
    "# scheduler_settings = (\"LambdaLR\", \"exp\", 2, False)    # Settings for the learning rate scheduler\n",
    "simplify_criteria = (\"DLs\", 0, 3, \"relative\") # The (criteria type, threshold, patience, compare_mode) upon which not satisfied, we break the current simplification and continue to the next layer/model\n",
    "optim_type = (\"adam\", 5e-3)\n",
    "optim_domain_type = (\"adam\", 1e-3)\n",
    "optim_autoencoder_type = (\"adam\", 1e-5, 1e-1) # optim_type, lr, loss_scale\n",
    "reg_mode = \"L1\"\n",
    "reg_amp = 1e-8\n",
    "reg_smooth = None\n",
    "reg_domain_mode = \"L1\"\n",
    "reg_domain_amp = 1e-5\n",
    "batch_size = 10000\n",
    "loss_core = \"DLs\"\n",
    "loss_order = -1\n",
    "loss_decay_scale = None\n",
    "is_mse_decay = False\n",
    "num_examples = 20000\n",
    "epochs = 10000\n",
    "iter_to_saturation = int(epochs / 2)\n",
    "MDL_mode = \"both\"\n",
    "date_time = \"{0}-{1}\".format(datetime.datetime.now().month, datetime.datetime.now().day)\n",
    "seed = 0\n",
    "array_id = \"0\"\n",
    "\n",
    "loss_balance_model_influence = False\n",
    "loss_success_threshold = 1e-4   # MSE level you regard as success\n",
    "theory_add_mse_threshold = 0.05 # MSE level below which you will add to the theory hub\n",
    "theory_remove_fraction_threshold = 0.005  # Fraction threshold below which you will remove a theory after each stage of training.\n",
    "matching_numerical_tolerance = 2e-4 # The tolerance below which you regard the numerical coefficient matches.\n",
    "matching_snapped_tolerance = 1e-9   # The tolerance below which you regard the snapped coefficient matches.\n",
    "max_trial_times = 1         # Maximum number of trial times before going on to next target (DEFAULT=1)\n",
    "is_simplify_model = True    # Whether to perform simplification of theory models\n",
    "is_simplify_domain = False  # Whether to perform simplification of theory domains\n",
    "record_mode = 2             # Record data mode. Choose from 0 (minimal recording), 1, 2 (record everything)\n",
    "show_3D_plot = False\n",
    "show_vs = False\n",
    "\n",
    "big_domain_dict = [(key, [1, 2]) for key in get_mystery([20000, 30000, 40000, 50000], range(4, 7), range(11))] + \\\n",
    "                  [(key, [1, 2]) for key in get_mystery([40000, 50000], [20], range(11))] + \\\n",
    "                  [(key, [1, 2, 3]) for key in get_mystery([20000, 30000, 40000, 50000], range(7, 10), range(11))] + \\\n",
    "                  [(key, [1, 2, 3, 4]) for key in get_mystery([20000, 30000, 40000, 50000], [10], range(11))]\n",
    "big_domain_dict = {key: item for key, item in big_domain_dict}\n",
    "########################################\n",
    "# Settings for double pendulum:\n",
    "is_pendulum = False\n",
    "is_pendulum = get_args(is_pendulum, 27, \"bool\")\n",
    "if is_pendulum:  \n",
    "    num_layers = 5\n",
    "    pred_nets_neurons = 160\n",
    "    domain_net_neurons = 16\n",
    "    reg_domain_amp = 1e-7\n",
    "    show_3D_plot = False\n",
    "    pred_nets_activation = \"tanh\"\n",
    "    csv_filename_list = ['mysteryt_px_10', 'mysteryt_px_11', 'mysteryt_px_12', 'mysteryt_px_13', 'mysteryt_px_20', 'mysteryt_px_21', 'mysteryt_px_23']\n",
    "    is_simplify_model = False\n",
    "    is_simplify_domain = False\n",
    "    add_theory_loss_threshold = 1e-4\n",
    "    num_theories_init = 2\n",
    "    add_theory_limit = 2\n",
    "    is_xv = True\n",
    "    record_mode = 1\n",
    "#     reg_smooth = (0.1, 2, 10, 1e-2, 1)\n",
    "    if is_xv:\n",
    "        csv_filename_list = ['mysteryt_pxv_10', 'mysteryt_pxv_11', 'mysteryt_pxv_12', 'mysteryt_pxv_13', 'mysteryt_pxv_20', 'mysteryt_pxv_21', 'mysteryt_pxv_23']\n",
    "        num_output_dims = 4\n",
    "        num_input_steps = 1\n",
    "        add_theory_loss_threshold = 3e-4\n",
    "    for key in csv_filename_list:\n",
    "        big_domain_dict[key] = [0, 1]\n",
    "########################################\n",
    "\n",
    "########################################\n",
    "# Settings for data_format = \"images\":\n",
    "if data_format == \"images\":\n",
    "    batch_size = 100\n",
    "    epochs = 10000\n",
    "    loss_core = \"mse\"\n",
    "    add_theory_quota = 0\n",
    "    is_simplify_model = False\n",
    "    is_simplify_domain = False\n",
    "########################################\n",
    "\n",
    "########################################\n",
    "# Settings for Lagrangian:\n",
    "if is_Lagrangian:\n",
    "    num_input_steps = 3\n",
    "    is_simplify_model = False\n",
    "########################################\n",
    "\n",
    "exp_id = get_args(exp_id, 1)\n",
    "env_source = get_args(env_source, 2)\n",
    "exp_mode = get_args(exp_mode, 3)\n",
    "num_theories_init = get_args(num_theories_init, 4, \"int\")\n",
    "pred_nets_neurons = get_args(pred_nets_neurons, 5, \"int\")\n",
    "pred_nets_activation = get_args(pred_nets_activation, 6)\n",
    "domain_net_neurons = get_args(domain_net_neurons, 7, \"int\")\n",
    "domain_pred_mode = get_args(domain_pred_mode, 8)\n",
    "mse_amp = get_args(mse_amp, 9, \"float\")\n",
    "simplify_criteria = get_args(simplify_criteria, 10, \"tuple\")\n",
    "scheduler_settings = get_args(scheduler_settings, 11, \"tuple\")\n",
    "optim_type = get_args(optim_type, 12, \"tuple\")\n",
    "optim_domain_type = get_args(optim_domain_type, 13, \"tuple\")\n",
    "reg_amp = get_args(reg_amp, 14, \"float\")\n",
    "reg_domain_amp = get_args(reg_domain_amp, 15, \"float\")\n",
    "batch_size = get_args(batch_size, 16, \"int\")\n",
    "loss_core = get_args(loss_core, 17)\n",
    "loss_order = get_args(loss_order, 18)\n",
    "loss_decay_scale = get_args(loss_decay_scale, 19, \"eval\")\n",
    "is_mse_decay = get_args(is_mse_decay, 20, \"bool\")\n",
    "loss_balance_model_influence = get_args(loss_balance_model_influence, 21, \"bool\")\n",
    "num_examples = get_args(num_examples, 22, \"int\")\n",
    "iter_to_saturation = get_args(iter_to_saturation, 23, \"eval\")\n",
    "MDL_mode = get_args(MDL_mode, 24)\n",
    "num_output_dims = get_args(num_output_dims, 25, \"int\")\n",
    "num_layers = get_args(num_layers, 26, \"int\")\n",
    "date_time = get_args(date_time, 28)\n",
    "seed = get_args(seed, 29, \"int\")\n",
    "array_id = get_args(array_id, 30, \"int\")\n",
    "is_batch = False\n",
    "array_id_core = array_id\n",
    "if isinstance(env_source, str) and \"mystery\" in env_source:\n",
    "    csv_filename_list = [env_source]\n",
    "    env_source = \"file\"\n",
    "    array_id = seed\n",
    "    load_previous = False\n",
    "    is_batch = True\n",
    "\n",
    "if exp_mode == \"continuous\":\n",
    "    is_propose_models = True\n",
    "elif exp_mode == \"newb\":\n",
    "    is_propose_models = False\n",
    "elif exp_mode == \"base\":\n",
    "    is_propose_models = False\n",
    "    num_theories_init = 1\n",
    "    pred_nets_activation = \"leakyRelu\" if pred_nets_activation == \"linear\" else pred_nets_activation\n",
    "    pred_nets_neurons *= 2\n",
    "    epochs *= 5\n",
    "    add_theory_quota = 0\n",
    "    is_simplify_model = False\n",
    "    is_simplify_domain = False\n",
    "    loss_core = \"mse\"\n",
    "    MDL_mode = \"None\"\n",
    "else:\n",
    "    raise\n",
    "batch_size = min(batch_size, num_examples)\n",
    "save_image = True\n",
    "render = False\n",
    "loss_floor = 1e-12\n",
    "simplify_lr = 1e-6\n",
    "simplify_epochs = 40000\n",
    "simplify_patience = 200\n",
    "target_symbolic_expressions = {}\n",
    "view_init = (10, 190)           # Angle you want to view the 3D plots\n",
    "\n",
    "\n",
    "if iter_to_saturation is not None:\n",
    "    iter_to_saturation = int(iter_to_saturation)\n",
    "    reg_multiplier = np.linspace(0, 1, iter_to_saturation) ** 2\n",
    "else:\n",
    "    reg_multiplier = None\n",
    "\n",
    "if epochs <= 3000:\n",
    "    change_interval = 2\n",
    "    record_interval = 5\n",
    "else:\n",
    "    change_interval = 10\n",
    "    record_interval = 5\n",
    "\n",
    "np.random.seed(seed)\n",
    "random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "dirname = theory_PATH + \"/{0}_{1}/\".format(exp_id, date_time)\n",
    "filename = dirname + \"{0}_{1}_num_{2}_pred_{3}_{4}_dom_{5}_{6}_mse_{7}_sim_{8}_optim_{10}_{11}_reg_{12}_{13}_batch_{14}_core_{15}_order_{16}_lossd_{17}_{18}_infl_{19}_#_{20}_mul_{21}_MDL_{22}_{23}D_{24}L_id_{25}_{26}.p\".format(\n",
    "            env_source, exp_mode, num_theories_init, pred_nets_neurons, pred_nets_activation, domain_net_neurons, domain_pred_mode, mse_amp,\n",
    "            to_string(simplify_criteria), to_string(scheduler_settings, num_strings = 5), to_string(optim_type), to_string(optim_domain_type), reg_amp, reg_domain_amp, batch_size, loss_core, \n",
    "            loss_order, loss_decay_scale, is_mse_decay, loss_balance_model_influence, num_examples, iter_to_saturation, MDL_mode, num_output_dims, num_layers, seed, array_id,\n",
    ")\n",
    "if is_batch:\n",
    "    filename_batch = filename[:-2] + \"_model/\" + \"{0}_{1}_num_{2}_pred_{3}_{4}_dom_{5}_{6}_mse_{7}_sim_{8}_optim_{10}_{11}_reg_{12}_{13}_batch_{14}_core_{15}_order_{16}_lossd_{17}_{18}_infl_{19}_#_{20}_mul_{21}_MDL_{22}_{23}D_{24}L_id_{25}_{26}.p\".format(\n",
    "            csv_filename_list[0], exp_mode, num_theories_init, pred_nets_neurons, pred_nets_activation, domain_net_neurons, domain_pred_mode, mse_amp,\n",
    "            to_string(simplify_criteria), to_string(scheduler_settings, num_strings = 5), to_string(optim_type), to_string(optim_domain_type), reg_amp, reg_domain_amp, batch_size, loss_core, \n",
    "            loss_order, loss_decay_scale, is_mse_decay, loss_balance_model_influence, num_examples, iter_to_saturation, MDL_mode, num_output_dims, num_layers, seed, array_id_core,\n",
    "    )\n",
    "    make_dir(filename_batch)\n",
    "make_dir(filename)\n",
    "make_dir(filename[:-2] + \"_mys/file\")\n",
    "if save_image:\n",
    "    make_dir(filename[:-2] + \"/images\")\n",
    "\n",
    "# Load theory_hub:\n",
    "filename_hub = filename[:-2] + \"_hub.p\"\n",
    "print(filename)\n",
    "\n",
    "# Load or initialize theory hub:\n",
    "load_augmenting_hub = False\n",
    "if load_augmenting_hub:\n",
    "    is_propose_models = True\n",
    "    dirname_augmenting_hub = theory_PATH + \"/{0}_{1}/\".format(\"transfer-hit5.5\", \"10-9\")\n",
    "    hub_candi1 = \"{0}_{1}_num_{2}_pred_{3}_{4}_dom_{5}_{6}_mse_{7}_sim_{8}_optim_{10}_{11}_reg_{12}_{13}\".format(\n",
    "                env_source, \"continuous\", num_theories_init, pred_nets_neurons, pred_nets_activation, domain_net_neurons, domain_pred_mode, mse_amp,\n",
    "                to_string(simplify_criteria), to_string(scheduler_settings, num_strings = 5), to_string(optim_type), to_string(optim_domain_type), reg_amp, reg_domain_amp, \n",
    "                \n",
    "    )\n",
    "    hub_candi2 = \"core_{0}_order_{1}_lossd_{2}_{3}_infl_{4}\".format(loss_core, loss_order, loss_decay_scale, is_mse_decay, loss_balance_model_influence)\n",
    "    hub_candi3 = \"mul_{0}_MDL_{1}_{2}D_{3}L_id_{4}\".format(iter_to_saturation, MDL_mode, num_output_dims, num_layers, seed, array_id)\n",
    "    filename_augmenting_hub = dirname_augmenting_hub + filter_filename(dirname_augmenting_hub, include = [hub_candi1, hub_candi2, hub_candi3, \"hub.p\"])[0]\n",
    "    theory_hub = load_model_dict_at_theory_hub(pickle.load(open(filename_augmenting_hub, \"rb\")), is_cuda = is_cuda)\n",
    "    pickle.dump(theory_hub.model_dict, open(filename_hub, \"wb\"))\n",
    "    print(\"load augmenting hub succeed: {0}\".format(filename_augmenting_hub))\n",
    "else:\n",
    "    if load_previous:\n",
    "        try:\n",
    "            theory_hub = load_model_dict_at_theory_hub(pickle.load(open(filename_hub, \"rb\")), is_cuda = is_cuda)\n",
    "        except:\n",
    "            theory_hub = Theory_Hub()\n",
    "            pickle.dump(theory_hub.model_dict, open(filename_hub, \"wb\"))\n",
    "    else:\n",
    "        theory_hub = Theory_Hub()\n",
    "        pickle.dump(theory_hub.model_dict, open(filename_hub, \"wb\"))\n",
    "\n",
    "\n",
    "def save_to_hub(pred_nets, domain_net, theory_hub, theory_type, theory_add_threshold, is_Lagrangian):\n",
    "    \"\"\"Save qualified theories to theory_hub\"\"\"\n",
    "    if load_previous:\n",
    "        theory_hub = load_model_dict_at_theory_hub(pickle.load(open(filename_hub, \"rb\")))\n",
    "    added_theory_info = theory_hub.add_theories(name = hub_theory_name if theory_type == \"neural\" else hub_theory_name + \"_simplified\",\n",
    "                                                pred_nets = pred_nets,\n",
    "                                                domain_net = domain_net,\n",
    "                                                dataset = dataset,\n",
    "                                                threshold = theory_add_threshold,\n",
    "                                                is_Lagrangian = is_Lagrangian,\n",
    "                                               )\n",
    "    if theory_type == \"neural\":\n",
    "        name = \"theory\"\n",
    "    elif theory_type == \"simplified\":\n",
    "        name = \"simplified_theory\"\n",
    "    else:\n",
    "        raise\n",
    "    info_dict[env_name][\"added_{0}_info\".format(name)] = added_theory_info\n",
    "    pickle.dump(theory_hub.model_dict, open(filename_hub, \"wb\"))\n",
    "\n",
    "def load_info_dict(info_dict, filename, trial_times = 1):\n",
    "    load_succeed = False\n",
    "    for i in range(trial_times):\n",
    "        try:\n",
    "            info_dict_updated = pickle.load(open(filename, \"rb\"))\n",
    "            print(\"Succesfully loaded {0}\".format(filename))\n",
    "            info_dict_updated.update(info_dict)\n",
    "            info_dict = info_dict_updated\n",
    "            load_succeed = True\n",
    "            break\n",
    "        except Exception as e:\n",
    "            print(\"Load trial {0}: \".format(i), e)\n",
    "            time.sleep(2 ** i * (0.8 + 0.4 * np.random.rand()))\n",
    "    return info_dict, load_succeed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory learning and simplification:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "info_dict = {}\n",
    "info_dict[\"array_id\"] = array_id\n",
    "info_dict[\"reg_smooth\"] = reg_smooth\n",
    "if load_previous:\n",
    "    info_dict, load_succeed = load_info_dict(info_dict, filename)\n",
    "    if not load_succeed:\n",
    "        pickle.dump(info_dict, open(filename, \"wb\"))\n",
    "\n",
    "###########################################################################################################################\n",
    "# Level I. Iterate over environments:\n",
    "###########################################################################################################################\n",
    "for env_name in csv_filename_list:\n",
    "    big_domain_ids = big_domain_dict[env_name] if env_name in big_domain_dict else None \n",
    "    # Dealing with specific environments:\n",
    "    print(\"\\n\" + \"=\" * 150 + \"\\nnew environment\\n\" + \"{0}\\n\".format(env_name) + \"=\" * 150 + \"\\n\\n\")\n",
    "    \n",
    "    # Check if the environment has been run. If so, skip:\n",
    "    if is_simplify_domain is True:\n",
    "        assert is_simplify_model is True\n",
    "        key_to_check = \"domain_net_simplified\"\n",
    "    elif is_simplify_model is True:\n",
    "        key_to_check = \"pred_nets_simplified\"\n",
    "    else:\n",
    "        key_to_check = \"pred_nets\"\n",
    "    if env_name not in info_dict:\n",
    "        info_dict[env_name] = {}\n",
    "    if key_to_check in info_dict[env_name]:\n",
    "        print(\"{0} already ran. Go to the next environment\".format(env_name))\n",
    "        continue\n",
    "    hub_theory_name = \"{0}_id_{1}\".format(env_name, array_id)\n",
    "    \n",
    "    # Load or generate dataset:\n",
    "    np.random.seed(seed)\n",
    "    random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    \n",
    "    # Get dataset from file:\n",
    "    dataset = get_dataset(filename=csv_dirname + env_name + \".csv\",\n",
    "                          num_output_dims=num_output_dims,\n",
    "                          is_classified=is_classified,\n",
    "                          num_input_steps=num_input_steps,\n",
    "                          forward_steps=forward_steps, \n",
    "                          num_examples=num_examples,\n",
    "                          is_Lagrangian=is_Lagrangian,\n",
    "                          is_cuda=is_cuda,\n",
    "                         )\n",
    "    if is_Lagrangian:\n",
    "        num_output_dims = 4\n",
    "        num_input_steps = 2\n",
    "    \n",
    "    ((X_train, y_train), (X_test, y_test), _), info = dataset\n",
    "    if num_output_dims <= 2 and \"true_domain_train\" in info:\n",
    "        plot_indi_domain(X_train, domain = info[\"true_domain_train\"], is_show = isplot, filename = filename[:-2] + \"/{0}_indi-domains.png\".format(env_name) if save_image else None)\n",
    "\n",
    "    # Setting up network and training hyperparameters:\n",
    "    # Structure for theory model:\n",
    "    if is_Lagrangian:\n",
    "        struct_param_pred = [[1, \"Symbolic_Layer\", {\"symbolic_expression\": \"g0 * x0 + g1 * x2\"}]]\n",
    "    else:\n",
    "        struct_param_pred = [[pred_nets_neurons, \"Simple_Layer\", {}]] * (num_layers - 1) + [[y_train.size(1), \"Simple_Layer\", {\"activation\": \"linear\"}]]\n",
    "    settings_pred = {\"activation\": pred_nets_activation}   # Default activation if not stipulated\n",
    "    \n",
    "    if data_format == \"images\":\n",
    "        X_train, y_train, X_test, y_test = to_pixels(X_train, y_train, X_test, y_test, size = (39, 39), xlim = (-2, 2), ylim = (-2, 2), radius = 3, shape = \"circle\", is_cuda = is_cuda)\n",
    "    \n",
    "    # Structure for theory domain:\n",
    "    struct_param_domain = [\n",
    "        [domain_net_neurons, \"Simple_Layer\", {}],\n",
    "        [domain_net_neurons, \"Simple_Layer\", {}],\n",
    "        [num_theories_init, \"Simple_Layer\", {\"activation\": \"linear\"}],\n",
    "    ]\n",
    "    settings_domain = {\"activation\": \"leakyRelu\"} # Default activation if not stipulated\n",
    "    loss_types = {\n",
    "                    \"pred-based_generalized-mean_{0}\".format(loss_order): {\"amp\": 1., \"decay_on\": True},\n",
    "                    \"pred-based_generalized-mean_1\": {\"amp\": mse_amp, \"decay_on\": is_mse_decay}, \n",
    "    }\n",
    "    \n",
    "    # Setting up regularization:\n",
    "    reg_dict = {\"pred_nets\": {\"weight\": reg_amp, \"bias\": reg_amp}}\n",
    "    reg_domain_dict = {\"domain_net\": {\"weight\": reg_domain_amp, \"bias\": reg_domain_amp}}\n",
    "    \n",
    "    # Settings for the simultaneous fitting of domain:\n",
    "    domain_fit_setting = {\"optim_domain_type\": optim_domain_type, \"reg_domain_dict\": reg_domain_dict, \"reg_domain_mode\": reg_domain_mode}\n",
    "    \n",
    "    \n",
    "    if data_format == \"images\":\n",
    "        enc_fully_size = 512\n",
    "        latent_size = 2\n",
    "        num_channels = 32\n",
    "        struct_param_encoder = [\n",
    "            [num_channels, \"Conv2d\", {\"kernel_size\": 3, \"stride\": 2}],\n",
    "            [num_channels, \"Conv2d\", {\"kernel_size\": 3, \"stride\": 2}],\n",
    "            [num_channels, \"Conv2d\", {\"kernel_size\": 3, \"stride\": 2}],\n",
    "            [latent_size, \"Simple_Layer\", {\"layer_input_size\": enc_fully_size}],\n",
    "        ]\n",
    "        struct_param_decoder = [\n",
    "            [enc_fully_size, \"Simple_Layer\", {\"layer_input_size\": latent_size}],\n",
    "            [num_channels, \"ConvTranspose2d\", {\"kernel_size\": 3, \"stride\": 2, \"layer_input_size\": (num_channels, 4, 4)}],\n",
    "            [num_channels, \"ConvTranspose2d\", {\"kernel_size\": 3, \"stride\": 2}],\n",
    "            [1, \"ConvTranspose2d\", {\"kernel_size\": 3, \"stride\": 2}],\n",
    "        ]\n",
    "        autoencoder = Conv_Autoencoder(\n",
    "            input_channels_encoder = 1,\n",
    "            input_channels_decoder = num_channels,\n",
    "            struct_param_encoder = struct_param_encoder,\n",
    "            struct_param_decoder = struct_param_decoder,\n",
    "            share_model_among_steps = True,\n",
    "            settings = {\"activation\": \"leakyReluFlat\"},\n",
    "            is_cuda = is_cuda,\n",
    "        )\n",
    "        print(\"Pre-training autoencoder:\")\n",
    "        train_simple(autoencoder, X_train, X_train, loss_type = \"mse\")\n",
    "        print(\"Pre-training autoencoder completes.\")\n",
    "        if isplot:\n",
    "            X_train_recons = autoencoder(X_train)\n",
    "            for i in np.random.randint(len(X_train), size = 4):\n",
    "                plot_matrices(torch.cat([X_train[i], X_train_recons[i]], 0), images_per_row = 4)\n",
    "        \n",
    "    else:\n",
    "        autoencoder = None\n",
    "\n",
    "    #########################################################################################################\n",
    "    # Level II.a. Iterate over number of trials:\n",
    "    #########################################################################################################\n",
    "    event_name = \"hit_until_loss_threshold\"\n",
    "    for trial_times in range(max_trial_times):\n",
    "        print(\"=\" * 120 + \"\\nTrial {0}\\n\".format(trial_times) + \"=\" * 120 + \"\\n\\n\")\n",
    "        # Each trial start with a different seed:\n",
    "        torch.manual_seed(seed + trial_times) \n",
    "        info_dict_single = {}\n",
    "\n",
    "        # Propose theory from theory_hub:\n",
    "        if load_previous:\n",
    "            theory_hub = load_model_dict_at_theory_hub(pickle.load(open(filename_hub, \"rb\")), is_cuda = is_cuda)\n",
    "        if is_propose_models:\n",
    "            print(\"From theory_hub:\")\n",
    "            proposed_theory_models, propose_evaluation = theory_hub.propose_theory_models(X_train, y_train, max_num_models = max(2, int(num_theories_init * 0.7)), types = [\"neural\"], is_Lagrangian = is_Lagrangian)\n",
    "        else:\n",
    "            proposed_theory_models = {}\n",
    "            propose_evaluation = {}\n",
    "        info_dict_proposed_models = {}\n",
    "        for name, theory_info in proposed_theory_models.items():\n",
    "            info_dict_proposed_models[name] = {}\n",
    "            for key, item in theory_info.items():\n",
    "                if key == \"theory_model\":\n",
    "                    info_dict_proposed_models[name][key] = item.model_dict\n",
    "                else:\n",
    "                    info_dict_proposed_models[name][key] = item\n",
    "        info_dict_single[\"proposed_theory_models\"] = info_dict_proposed_models\n",
    "        info_dict_single[\"proposed_theory_models_evaluation\"] = propose_evaluation\n",
    "\n",
    "\n",
    "        # Initialize Theory_Training:\n",
    "        T = Theory_Training(\n",
    "            num_theories = num_theories_init,\n",
    "            proposed_theory_models = proposed_theory_models,\n",
    "            input_size = info[\"input_size\"],\n",
    "            struct_param_pred = struct_param_pred,\n",
    "            struct_param_domain = struct_param_domain,\n",
    "            settings_pred = settings_pred,\n",
    "            settings_domain = settings_domain,\n",
    "            autoencoder = autoencoder,\n",
    "            loss_types = loss_types,\n",
    "            loss_core = loss_core,\n",
    "            loss_order = loss_order,\n",
    "            loss_balance_model_influence = loss_balance_model_influence,\n",
    "            is_Lagrangian = is_Lagrangian,\n",
    "            reg_multiplier = reg_multiplier,\n",
    "            is_cuda = is_cuda,\n",
    "        )\n",
    "\n",
    "        #################################################################################################\n",
    "        # Level III.a Iterative training of models and domains\n",
    "        #################################################################################################\n",
    "        print(\"\\n\" + \"=\" * 80)\n",
    "        print(\"{0}, trial times {1}, training model:\".format(env_name, trial_times))\n",
    "        print(\"=\" * 80 + \"\\n\")\n",
    "        loss_precision_floor_init = 10\n",
    "        T.set_loss_core(loss_core, loss_precision_floor_init)\n",
    "        data_record = T.iterative_train(\n",
    "            X_train,\n",
    "            y_train,\n",
    "            validation_data = (X_test, y_test),\n",
    "            optim_type = optim_type,\n",
    "            optim_autoencoder_type = optim_autoencoder_type,\n",
    "            reg_dict = reg_dict,\n",
    "            reg_mode = reg_mode,\n",
    "            reg_smooth = reg_smooth,\n",
    "            domain_fit_setting = domain_fit_setting,\n",
    "            forward_steps = forward_steps,\n",
    "            domain_pred_mode = domain_pred_mode,\n",
    "            scheduler_settings = scheduler_settings,\n",
    "            loss_order_decay = (lambda epoch: - epoch / float(loss_decay_scale)) if loss_decay_scale is not None else None,\n",
    "            view_init = view_init,\n",
    "            epochs = epochs,\n",
    "            patience = int(epochs / 50) if exp_mode != \"base\" else None,\n",
    "            batch_size = batch_size,\n",
    "            inspect_interval = int(epochs / 100) if is_pendulum else int(epochs / 10),\n",
    "            change_interval = change_interval,\n",
    "            record_interval = record_interval,\n",
    "            record_mode = record_mode,\n",
    "            isplot = isplot,\n",
    "            filename = filename[:-2] + \"/{0}_training-model\".format(env_name) if save_image else None,\n",
    "            raise_nan = False,\n",
    "            add_theory_quota = add_theory_quota,\n",
    "            add_theory_limit = add_theory_limit,\n",
    "            add_theory_criteria = add_theory_criteria,\n",
    "            add_theory_loss_threshold = add_theory_loss_threshold,\n",
    "            loss_floor = loss_floor,\n",
    "            true_domain_test = info[\"true_domain_test\"] if \"true_domain_test\" in info else None,\n",
    "            num_output_dims = num_output_dims,\n",
    "            prefix = \"{0}_train_model:\".format(env_name),\n",
    "            show_3D_plot = show_3D_plot,\n",
    "            show_vs = show_vs,\n",
    "            big_domain_ids = big_domain_ids,\n",
    "            fix_adaptive_precision_floor = True,\n",
    "        )\n",
    "        data_record[\"loss_precision_floor\"] = loss_precision_floor_init\n",
    "        info_dict_single[\"data_record\"] = deepcopy(data_record)\n",
    "        info_dict_single[\"pred_nets\"] = deepcopy(T.pred_nets.model_dict)\n",
    "        info_dict_single[\"removed_theories_0\"] = T.remove_theories_based_on_data(X_test, y_test, threshold = theory_remove_fraction_threshold)\n",
    "\n",
    "        # Perform MDL training after MSE training if stipulated:\n",
    "        if MDL_mode == \"both\":\n",
    "            if loss_core == \"mse\":\n",
    "                T.set_loss_core(\"DLs\")\n",
    "            if hasattr(T, \"optimizer\"):\n",
    "                delattr(T, \"optimizer\")\n",
    "            data_record_MDL1 = T.iterative_train_schedule(\n",
    "                X_train,\n",
    "                y_train,\n",
    "                validation_data = (X_test, y_test),\n",
    "                optim_type = (\"adam\", 1e-4),\n",
    "                reg_dict = reg_dict,\n",
    "                reg_mode = reg_mode,\n",
    "                reg_smooth = reg_smooth,\n",
    "                domain_fit_setting = domain_fit_setting,\n",
    "                forward_steps = forward_steps,\n",
    "                domain_pred_mode = domain_pred_mode,\n",
    "                scheduler_settings = scheduler_settings,\n",
    "                loss_order_decay = (lambda epoch: - epoch / float(loss_decay_scale)) if loss_decay_scale is not None else None,\n",
    "                view_init = view_init,\n",
    "                epochs = int(epochs / 2),\n",
    "                patience = int(epochs / 50) if exp_mode != \"base\" else None,\n",
    "                batch_size = batch_size,\n",
    "                inspect_interval = int(epochs / 100) if is_pendulum else int(epochs / 10),\n",
    "                change_interval = change_interval,\n",
    "                record_interval = record_interval,\n",
    "                record_mode = record_mode,\n",
    "                isplot = isplot,\n",
    "                filename = filename[:-2] + \"/{0}_training-modelDL1\".format(env_name) if save_image else None,\n",
    "                raise_nan = False,\n",
    "                add_theory_quota = add_theory_quota,\n",
    "                add_theory_limit = add_theory_limit,\n",
    "                add_theory_criteria = add_theory_criteria,\n",
    "                add_theory_loss_threshold = add_theory_loss_threshold,\n",
    "                theory_remove_fraction_threshold = theory_remove_fraction_threshold,\n",
    "                loss_floor = loss_floor,\n",
    "                true_domain_test = info[\"true_domain_test\"] if \"true_domain_test\" in info else None,\n",
    "                num_output_dims = num_output_dims,\n",
    "                prefix = \"{0}_train_modelDL1:\".format(env_name),\n",
    "                show_3D_plot = show_3D_plot,\n",
    "                show_vs = show_vs,\n",
    "                big_domain_ids = big_domain_ids,\n",
    "                num_phases = 4,\n",
    "            )\n",
    "            info_dict_single[\"data_record_MDL1\"] = deepcopy(data_record_MDL1)\n",
    "\n",
    "            T.domain_net_on = True\n",
    "            print(\"\\ndomain_net turned on, and each theory only optimize in its domain:\\n\")\n",
    "            data_record_MDL2 = T.iterative_train_schedule(\n",
    "                X_train,\n",
    "                y_train,\n",
    "                validation_data = (X_test, y_test),\n",
    "                optim_type = (\"adam\", 1e-4),\n",
    "                reg_dict = reg_dict,\n",
    "                reg_mode = reg_mode,\n",
    "                reg_smooth = reg_smooth,\n",
    "                domain_fit_setting = domain_fit_setting,\n",
    "                forward_steps = forward_steps,\n",
    "                domain_pred_mode = domain_pred_mode,\n",
    "                scheduler_settings = scheduler_settings,\n",
    "                loss_order_decay = (lambda epoch: - epoch / float(loss_decay_scale)) if loss_decay_scale is not None else None,\n",
    "                view_init = view_init,\n",
    "                epochs = int(epochs / 2),\n",
    "                patience = int(epochs / 50) if exp_mode != \"base\" else None,\n",
    "                batch_size = batch_size,\n",
    "                inspect_interval = int(epochs / 100) if is_pendulum else int(epochs / 10),\n",
    "                change_interval = change_interval,\n",
    "                record_interval = record_interval,\n",
    "                record_mode = record_mode,\n",
    "                isplot = isplot,\n",
    "                filename = filename[:-2] + \"/{0}_training-modelDL2\".format(env_name) if save_image else None,\n",
    "                raise_nan = False,\n",
    "                add_theory_quota = add_theory_quota,\n",
    "                add_theory_limit = add_theory_limit,\n",
    "                add_theory_criteria = add_theory_criteria,\n",
    "                add_theory_loss_threshold = add_theory_loss_threshold,\n",
    "                theory_remove_fraction_threshold = theory_remove_fraction_threshold,\n",
    "                loss_floor = loss_floor,\n",
    "                true_domain_test = info[\"true_domain_test\"] if \"true_domain_test\" in info else None,\n",
    "                num_output_dims = num_output_dims,\n",
    "                prefix = \"{0}_train_modelDL2:\".format(env_name),\n",
    "                show_3D_plot = show_3D_plot,\n",
    "                show_vs = show_vs,\n",
    "                big_domain_ids = big_domain_ids,\n",
    "                num_phases = 2,\n",
    "            )\n",
    "            info_dict_single[\"data_record_MDL2\"] = deepcopy(data_record_MDL2)\n",
    "            if hasattr(T, \"optimizer\"):\n",
    "                delattr(T, \"optimizer\")\n",
    "\n",
    "        # Remove theories whose fraction is below the remove_theshold:\n",
    "        fraction_list = T.get_fraction_list(X_test, y_test)\n",
    "        removed_theories = T.remove_theories_based_on_data(X_test, y_test, threshold = theory_remove_fraction_threshold)\n",
    "        info_dict_single[\"fraction_list\"] = fraction_list\n",
    "        info_dict_single[\"removed_theories\"] = removed_theories\n",
    "        info_dict_single[\"domain_net\"] = deepcopy(T.domain_net.model_dict)\n",
    "        \n",
    "        # Record pred_nets and domain_nets before simplification\n",
    "        pred_nets = deepcopy(T.pred_nets)\n",
    "        domain_net = deepcopy(T.domain_net)\n",
    "\n",
    "        #########################################################################################################\n",
    "        # Level III.b. Simplify theory models and/or domains:\n",
    "        #########################################################################################################\n",
    "        kwargs = {\"true_domain_test\": info[\"true_domain_test\"] if \"true_domain_test\" in info else None, \n",
    "                  \"big_domain_ids\": big_domain_ids,\n",
    "                  \"is_Lagrangian\": is_Lagrangian,\n",
    "                 }\n",
    "        if is_simplify_model:\n",
    "            print(\"\\n\" + \"=\" * 150)\n",
    "            print(\"{0}, {1}\".format(env_name, \"simplifying model:\"))\n",
    "            print(\"=\" * 150 + \"\\n\")\n",
    "            info_dict_single[\"data_record_simplification-model\"] = {}\n",
    "            record_data(info_dict_single[\"data_record_simplification-model\"], [T.get_losses(X_test, y_test, **kwargs), T.pred_nets.model_dict, \"before_simplification\"], [\"all_losses_dict\", \"pred_nets_model_dict\", \"event\"])\n",
    "\n",
    "            # Get prediction by theory domains:\n",
    "            valid_onehot = to_one_hot(T.domain_net_forward(X_train).max(1)[1], T.num_theories)\n",
    "            valid_onehot_test = to_one_hot(T.domain_net_forward(X_test).max(1)[1], T.num_theories)\n",
    "            validation_pred_nets = (X_test, y_test, valid_onehot_test)\n",
    "\n",
    "            # Collapse Layers:\n",
    "            loss_record = T.pred_nets.simplify(X_train, y_train, valid_onehot, \"collapse_layers\", simplify_criteria = simplify_criteria, validation_data = validation_pred_nets, is_Lagrangian = is_Lagrangian)\n",
    "            T.pred_nets.get_weights_bias(W_source = \"core\", b_source = \"core\", verbose = True)\n",
    "            record_data(info_dict_single[\"data_record_simplification-model\"], [T.get_losses(X_test, y_test, **kwargs), loss_record, T.pred_nets.model_dict, \"after_collapsing_layer\"], [\"all_losses_dict\", \"loss_record\", \"pred_nets_model_dict\", \"event\"])\n",
    "\n",
    "            # Local snapping:\n",
    "            loss_record = T.pred_nets.simplify(X_train, y_train, valid_onehot, \"local\", loss_type = \"DLs\", loss_precision_floor = T.loss_precision_floor, simplify_criteria = simplify_criteria, validation_data = validation_pred_nets, patience = simplify_patience, lr = simplify_lr, epochs = simplify_epochs, is_Lagrangian = is_Lagrangian, verbose = 2)\n",
    "            record_data(info_dict_single[\"data_record_simplification-model\"], [T.get_losses(X_test, y_test, **kwargs), loss_record, T.pred_nets.model_dict, \"after_local_snapping\"], [\"all_losses_dict\", \"loss_record\", \"pred_nets_model_dict\", \"event\"])\n",
    "\n",
    "            # Integer snapping:\n",
    "            loss_record = T.pred_nets.simplify(X_train, y_train, valid_onehot, \"snap\", snap_mode = \"integer\", loss_type = \"DLs\", loss_precision_floor = T.loss_precision_floor, simplify_criteria = simplify_criteria, validation_data = validation_pred_nets, patience = simplify_patience, lr = simplify_lr, epochs = simplify_epochs, is_Lagrangian = is_Lagrangian, verbose = 2)\n",
    "            record_data(info_dict_single[\"data_record_simplification-model\"], [T.get_losses(X_test, y_test, **kwargs), loss_record, T.pred_nets.model_dict, \"after_integer_snapping\"], [\"all_losses_dict\", \"loss_record\", \"pred_nets_model_dict\", \"event\"])\n",
    "\n",
    "            # Rational snapping:\n",
    "            loss_record = T.pred_nets.simplify(X_train, y_train, valid_onehot, \"snap\", snap_mode = \"rational\", loss_type = \"DLs\", loss_precision_floor = T.loss_precision_floor, simplify_criteria = simplify_criteria, validation_data = validation_pred_nets, patience = simplify_patience, lr = simplify_lr, epochs = simplify_epochs, is_Lagrangian = is_Lagrangian, verbose = 2)\n",
    "            record_data(info_dict_single[\"data_record_simplification-model\"], [T.get_losses(X_test, y_test, **kwargs), loss_record, T.pred_nets.model_dict, \"after_rational_snapping\"], [\"all_losses_dict\", \"loss_record\", \"pred_nets_model_dict\", \"event\"])\n",
    "            \n",
    "            # To symbolic:\n",
    "            loss_record = T.pred_nets.simplify(X_train, y_train, valid_onehot, \"to_symbolic\", simplify_criteria = simplify_criteria, validation_data = validation_pred_nets, is_Lagrangian = is_Lagrangian)\n",
    "            T.pred_nets.get_sympy_expression()\n",
    "            record_data(info_dict_single[\"data_record_simplification-model\"], [T.get_losses(X_test, y_test, **kwargs), loss_record, T.pred_nets.model_dict, \"after_to_symbolic_and_snapping\"], [\"all_losses_dict\", \"loss_record\", \"pred_nets_model_dict\", \"event\"])\n",
    "\n",
    "            # Saving and show expression:\n",
    "            info_dict_single[\"pred_nets_simplified\"] = deepcopy(T.pred_nets.model_dict)\n",
    "            print()\n",
    "            print(\"MSE curve for simplification of the theory_model:\")\n",
    "            if isplot:\n",
    "                try:\n",
    "                    plot_loss_record(info_dict_single['data_record_simplification-model'], T.pred_nets.num_models)\n",
    "                except:\n",
    "                    pass\n",
    "            if save_image:\n",
    "                T.plot(X_test, y_test, true_domain = info[\"true_domain_test\"] if \"true_domain_test\" in info else None, \n",
    "                       view_init = view_init, figsize = (10, 8), filename = filename[:-2] + \"/{0}_model-simplified\".format(env_name) if save_image else None, is_show = isplot)\n",
    "\n",
    "            # Simplifying domain:\n",
    "            if is_simplify_domain:\n",
    "                print(\"\\n\" + \"=\" * 150)\n",
    "                print(\"{0}, {1}\".format(env_name, \"simplifying domain:\"))\n",
    "                print(\"=\" * 150 + \"\\n\")\n",
    "                info_dict_single[\"data_record_simplification-domain-model\"] = {}\n",
    "                record_data(info_dict_single[\"data_record_simplification-domain-model\"], [T.get_losses(X_test, y_test, **kwargs), T.domain_net.model_dict, \"before_domain_simplification\"], [\"all_losses_dict\", \"domain_net_model_dict\", \"event\"])\n",
    "\n",
    "                best_theory_idx = get_best_model_idx(T.net_dict, X_train, y_train, loss_fun_cumu = T.loss_fun_cumu)\n",
    "                loss_record = T.domain_net.simplify(X_train, best_theory_idx, mode = \"to_symbolic\", loss_type = \"cross-entropy\", simplify_criteria = simplify_criteria)\n",
    "                record_data(info_dict_single[\"data_record_simplification-domain-model\"], [T.get_losses(X_test, y_test, **kwargs), loss_record, T.domain_net.model_dict, \"after_to_symbolic\"], [\"all_losses_dict\", \"loss_record\", \"domain_net_model_dict\", \"event\"])\n",
    "\n",
    "                loss_record = T.domain_net.simplify(X_train, best_theory_idx, mode = \"pair_snap\", loss_type = \"cross-entropy\", snap_mode = \"integer\", simplify_criteria = simplify_criteria)\n",
    "                record_data(info_dict_single[\"data_record_simplification-domain-model\"], [T.get_losses(X_test, y_test, **kwargs), loss_record, T.domain_net.model_dict, \"after_pair_snap\"], [\"all_losses_dict\", \"loss_record\", \"domain_net_model_dict\", \"event\"])\n",
    "                info_dict_single[\"domain_net_simplified_pair_snap\"] = deepcopy(T.domain_net.model_dict)\n",
    "\n",
    "                loss_record = T.domain_net.simplify(X_train, best_theory_idx, mode = \"activation_snap\", loss_type = \"cross-entropy\", activation_source = \"sigmoid\", \n",
    "                                                    activation_target = \"heaviside\", epochs = 10000, simplify_criteria = simplify_criteria)\n",
    "                record_data(info_dict_single[\"data_record_simplification-domain-model\"], [T.get_losses(X_test, y_test, **kwargs), loss_record, T.domain_net.model_dict, \"after_activation_snap\"], [\"all_losses_dict\", \"loss_record\", \"domain_net_model_dict\", \"event\"])\n",
    "                info_dict_single[\"domain_net_simplified_activation_snap\"] = deepcopy(T.domain_net.model_dict)\n",
    "\n",
    "                loss_record = T.domain_net.simplify(X_train, best_theory_idx, mode = \"pair_snap\", loss_type = \"cross-entropy\", snap_mode = \"integer\", simplify_criteria = simplify_criteria)\n",
    "                record_data(info_dict_single[\"data_record_simplification-domain-model\"], [T.get_losses(X_test, y_test, **kwargs), loss_record, T.domain_net.model_dict, \"after_second_pair_snap\"], [\"all_losses_dict\", \"loss_record\", \"domain_net_model_dict\", \"event\"])\n",
    "                info_dict_single[\"domain_net_simplified_final\"] = deepcopy(T.domain_net.model_dict)\n",
    "                print(\"\\nSimplified domain expressions:\")\n",
    "                T.domain_net.show_expression()\n",
    "                print()\n",
    "                if save_image:\n",
    "                    T.plot(X_test, y_test, true_domain = info[\"true_domain_test\"] if \"true_domain_test\" in info else None, \n",
    "                           view_init = view_init, figsize = (10, 8), filename = filename[:-2] + \"/{0}_domain-simplified\".format(env_name) if save_image else None, is_show = isplot)\n",
    "\n",
    "            print(\"\\n\\n\")\n",
    "\n",
    "        if isplot:\n",
    "            try:\n",
    "                process_loss(info_dict_single, loss_core = loss_core)\n",
    "            except:\n",
    "                pass\n",
    "\n",
    "        # Record and determine whether to run another trial:\n",
    "        loss_dict = T.get_losses(X_test, y_test)\n",
    "        loss_evaluation = loss_dict[\"mse_with_domain\"] if exp_mode != \"base\" else loss_dict[\"mse_without_domain\"]\n",
    "        record_data(info_dict[env_name], [loss_evaluation, info_dict_single, event_name], [\"loss_evaluation\", \"info_dict_single\", \"event\"])\n",
    "#         info_dict, _ = load_info_dict(info_dict, filename)\n",
    "        print(\"#\" * 70 + \"\\nEvaluation for current trial for {0}:\\n\".format(env_name) + \"#\" * 70)\n",
    "        if loss_evaluation < loss_success_threshold:\n",
    "            if exp_mode in [\"continuous\", \"newb\"]:\n",
    "                event_name = \"hit_until_symbolic_matching\"\n",
    "                if env_name not in target_symbolic_expressions or target_symbolic_expressions[env_name] is None or is_simplify_model is False:\n",
    "                    print(\"When the target_symbolic_expression for the corresponding env_name is None or 'is_simplify_model' is False, cannot compare symbolic matching.\")\n",
    "                    print(\"\\nThe loss is below the success threshold {0}. Break.\\n\".format(loss_success_threshold))\n",
    "                    break\n",
    "                else:\n",
    "                    is_symbolic_matching, num_matches = check_expression_matching(T.pred_nets.get_sympy_expression(), \n",
    "                                                                                  target_symbolic_expressions[env_name] if env_name in target_symbolic_expressions else [standardize(\"0\")], \n",
    "                                                                                  tolerance = matching_numerical_tolerance, \n",
    "                                                                                  snapped_tolerance = matching_snapped_tolerance)\n",
    "                    if is_symbolic_matching:\n",
    "                        print(\"\\nThe symbolic expressions fully match the target expressions. Break.\\n\")\n",
    "                        break\n",
    "                    else:\n",
    "                        print(\"\\n{0} out of {1} expressions match the target expression {2}. Do another trial.\\n\".format(num_matches, len(T.pred_nets.get_sympy_expression()), target_symbolic_expressions[env_name]))\n",
    "            elif exp_mode == \"base\":\n",
    "                print(\"\\nThe loss is below the success threshold {0}. Break.\\n\".format(loss_success_threshold))\n",
    "                break\n",
    "            else:\n",
    "                raise Exception(\"exp_mode {0} not recognized!\".format(exp_mode))\n",
    "        else:\n",
    "            print(\"\\nThe loss does not get down to {0}, Do another trial.\".format(loss_success_threshold))\n",
    "        print(\"#\" * 70 + \"\\n\")\n",
    "\n",
    "    # Record the last trial (either succeed or at the max_trial):\n",
    "    export_csv_with_domain_prediction(env_name=env_name,\n",
    "                                      domain_net=T.domain_net,\n",
    "                                      num_output_dims=num_output_dims,\n",
    "                                      num_input_steps=num_input_steps,\n",
    "                                      csv_dirname=csv_dirname,\n",
    "                                      write_dirname=filename[:-2] + \"_mys/\",\n",
    "                                      is_Lagrangian=is_Lagrangian,\n",
    "                                      is_cuda=is_cuda,\n",
    "                                     )\n",
    "    info_dict[env_name][\"pred_nets\"] = pred_nets.model_dict\n",
    "    info_dict[env_name][\"domain_net\"] = domain_net.model_dict\n",
    "    if hasattr(T, \"autoencoder\"):\n",
    "        info_dict[env_name][\"autoencoder\"] = T.autoencoder.model_dict\n",
    "    if is_simplify_model:\n",
    "        info_dict[env_name][\"pred_nets_simplified\"] = T.pred_nets.model_dict\n",
    "    if is_simplify_domain:\n",
    "        info_dict[env_name][\"domain_net_simplified\"] = T.domain_net.model_dict\n",
    "    # Add theory to the theory hub:\n",
    "    save_to_hub(pred_nets, domain_net, theory_hub, theory_type = \"neural\", theory_add_threshold = theory_add_mse_threshold, is_Lagrangian = is_Lagrangian)\n",
    "    save_to_hub(T.pred_nets, T.domain_net, theory_hub, theory_type = \"simplified\", theory_add_threshold = theory_add_mse_threshold, is_Lagrangian = is_Lagrangian)\n",
    "\n",
    "    # Save files:\n",
    "    if not is_batch:\n",
    "        pickle.dump(info_dict, open(filename, \"wb\"))\n",
    "    else:\n",
    "        pickle.dump(info_dict, open(filename_batch, \"wb\"))\n",
    "\n",
    "    # Write the learned expression to csv_file:\n",
    "    coeff_learned_list, is_snapped_list = get_coeff_learned_list(T.pred_nets)\n",
    "    open(filename[:-2] + \"_mys/\" + env_name + \"_learned.csv\", 'w').close()\n",
    "    open(filename[:-2] + \"_mys/\" + env_name + \"_snapped.csv\", 'w').close()\n",
    "    with open(filename[:-2] + \"_mys/\" + env_name + \"_learned.csv\", \"a\") as f:\n",
    "        for coeff in coeff_learned_list:\n",
    "            if len(coeff) > 0:\n",
    "                line = \",\".join([str(element) for element in coeff]) + '\\n'\n",
    "                f.write(line)\n",
    "    with open(filename[:-2] + \"_mys/\" + env_name + \"_snapped.csv\", \"a\") as f:\n",
    "        for is_snapped in is_snapped_list:\n",
    "            if len(is_snapped) > 0:\n",
    "                line = \",\".join([str(element) for element in is_snapped]) + '\\n'\n",
    "                f.write(line)\n",
    "\n",
    "    if isplot:\n",
    "        try:\n",
    "            process_loss(info_dict[env_name][\"info_dict_single\"][-1], loss_core = loss_core)\n",
    "        except Exception as e:\n",
    "            print(e)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
